from dataset_create import TraficSignDataSet
from cnn_create import CNN
from torchvision import transforms
import torch
import torch.nn as nn
import os

def train():

    num_epochs = 5
    num_classes = 43
    batch_size = 50
    learning_rate = 0.001

    # Import dataset
    train_dataset = TraficSignDataSet(need_train=True, transform=transforms.ToTensor())

    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               shuffle=True)

    # Create the CNN model
    net = CNN(num_classes)
    print(net)
"""
    # Loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)

    # Train the model
    total_step = len(train_loader)

    for epoch in range(num_epochs):
        for i, data in enumerate(train_loader):
            # Forward pass
            image, label = data['image'], data['label']
            output = net(image)
            loss = criterion(output, label)
            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if (i + 1) % 50 == 0:
                print('Epoch [{}, {}], Step [{}, {}], Loss: {:.4f}'
                      .format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))

    print('Train Finish.')
    # Save the model checkpoint
    torch.save(net.state_dict(), 'net.ckpt')
    print('Net model has been saved.')
"""
if __name__ == '__main__':
    ans = input('Are you sure to train the model ? [y/n]')
    if ans is 'n':
        os._exit(0)

    pid = os.fork()
    if pid != 0:
        os._exit(0)
    else:
        train()
